"""
    CNNVD 漏洞查询
    https://www.cnnvd.org.cn/home/loophole
"""
import pprint
import requests
import openpyxl
import json
from my_fake_useragent import UserAgent
import os
from bridge.context import ContextType
import hashlib
import time
from common.log import logger
import requests
from bridge.reply import Reply, ReplyType
from plugins import *

hazardLevel ={
    0:"未知❓",
    1:"超危❗❗❗",
    2:"高危❗❗",
    3:"中危❗",
    4:"低危⚠️"
}

def search_cnnvd(number,e_context:EventContext):
    """
    number : CNNVD编号 例如 CNNVD-202311-1534
    """
    #
    content = f"🔍【{number}】\n"
    # 最大二次请求条数
    max_data = 2
    logger.info(f"🔍【{number}】")

    headers = {
        "Accept": "application/json, text/plain, */*",
        "Accept-Language": "zh-CN,zh;q=0.9",
        "Cache-Control": "no-cache",
        "Connection": "keep-alive",
        "Content-Type": "application/json;charset=UTF-8",
        "Origin": "https://www.cnnvd.org.cn",
        "Pragma": "no-cache",
        "Referer": "https://www.cnnvd.org.cn/home/loophole",
        "Sec-Fetch-Dest": "empty",
        "Sec-Fetch-Mode": "cors",
        "Sec-Fetch-Site": "same-origin",
        "User-Agent": UserAgent().random(),
        "sec-ch-ua": "^\\^Google",
        "sec-ch-ua-mobile": "?0",
        "sec-ch-ua-platform": "^\\^Windows^^",
    }
    url = "https://www.cnnvd.org.cn/web/homePage/cnnvdVulList"
    data = '{{"pageIndex":1,"pageSize":{},"keyword":"{}","hazardLevel":"","vulType":"","vendor":"","product":"","dateType":""}}'.format(
        max_data, number)

    response = requests.post(url, headers=headers, data=data)
    """
    {
    'code': 200,
    'data': {'pageIndex': 1,
              'pageSize': 10,
              'records': [{'cnnvdCode': 'CNNVD-202311-1534',
                           'createTime': '2023-11-17',
                           'cveCode': 'CVE-2023-6178',
                           'hazardLevel': 0,
                           'id': 'b53f7d7e933e4876b5e8452153719738',
                           'publishTime': '2023-11-17',
                           'typeName': None,
                           'updateTime': '2023-11-19',
                           'vulName': 'Nessus 安全漏洞',
                           'vulType': '0'}],
              'total': 1},
    'message': '操作成功',
    'success': True,
    'time': '2023-11-20 15:54:26'
    }
    """
    try:
        response_data = response.json()
        response_status = response_data.get('code',400)
        print(response_data)
        if response_status == 200:
            total = response_data.get('data').get('total')
            print(f"总共{total}条")
            logger.info(f"  (共查询到{total}条漏洞信息)")
            # 如果条数为0
            if total == 0:
                content += "未查询到相关漏洞信息"

            # 如果条数小于等于设定的max_data
            elif total <= max_data:
                content += f"  (共查询到{total}条漏洞信息)\n"
                records = response_data.get('data').get('records')
                for record in records:
                    loophole_id = record.get('id')
                    cnnvd_code = record.get('cnnvdCode')
                    content += get_loophole(loophole_id,cnnvd_code=cnnvd_code)

            # 如果条数大于设定的max_data
            # 则发送Excel
            elif total > max_data:
                _send_info(e_context=e_context,content=f"🔍【{number}】相关信息共{total}条，已整理为Excel概览")
                frequency = total // 50
                frequency = 1 if frequency == 0 else frequency

                print(f"需要请求{frequency}次")
                result_data = []
                for page in range(1,frequency+1):
                    print(f"第{page}次请求")
                    url = "https://www.cnnvd.org.cn/web/homePage/cnnvdVulList"
                    data = '{{"pageIndex":{},"pageSize":{},"keyword":"{}","hazardLevel":"","vulType":"","vendor":"","product":"","dateType":""}}'.format(
                        page,50, number)
                    result_data += requests.post(url, headers=headers, data=data).json().get('data').get('records')
                file_path = "tmp"

                if not os.path.exists(file_path):
                    os.makedirs(file_path)

                file_name = hashlib.md5(str(time.time()).encode()).hexdigest() + ".xlsx"
                file_path = os.path.join(file_path, file_name)

                print(file_path)
                # 创建一个新的工作簿
                workbook = openpyxl.Workbook()
                # 选择默认的工作表
                sheet = workbook.active
                sheet.append(['漏洞编号','漏洞等级','漏洞名称','收录时间','更新时间','查询id',])

                print(len(result_data))
                for record in result_data:
                    data = [
                        record.get('cnnvdCode'),
                        hazardLevel[record.get('hazardLevel')],
                        record.get('vulName'),
                        record.get('createTime'),
                        record.get('updateTime'),
                        record.get('id'),
                    ]
                    print(data)
                    sheet.append(data)
                sheet.append([])
                sheet.append([f"信息量大于{max_data}条"])
                sheet.append([f"已自动为您生成【{number}】相关简报"])
                # 保存工作簿
                workbook.save(file_path)
                reply = Reply()
                # 发送文件
                reply.type = ReplyType.FILE
                reply.content = file_path
                logger.info(f"查询结果：{file_path}")
                e_context["reply"] = reply
                e_context.action = EventAction.BREAK_PASS
                # 返回：暂时无用
                return file_path,reply

    except Exception as e:
        print(e)
        logger.info(f"查询漏洞发生错误:{e}")
        content += f"发生错误:{e}\n"

    reply = Reply()
    reply.type = ReplyType.TEXT
    reply.content = content
    logger.info(reply.content)
    e_context["reply"] = reply
    e_context.action = EventAction.BREAK_PASS
    # 返回：暂时无用
    return content

# 进一步查询漏洞详细信息
def get_loophole(loophole_id,cnnvd_code="CNNVD-202311-299"):
    content = ""
    headers = {
        "Accept": "application/json, text/plain, */*",
        "Accept-Language": "zh-CN,zh;q=0.9",
        "Cache-Control": "no-cache",
        "Connection": "keep-alive",
        "Content-Type": "application/json;charset=UTF-8",
        "Origin": "https://www.cnnvd.org.cn",
        "Pragma": "no-cache",
        "Referer": "https://www.cnnvd.org.cn/home/loophole",
        "Sec-Fetch-Dest": "empty",
        "Sec-Fetch-Mode": "cors",
        "Sec-Fetch-Site": "same-origin",
        "User-Agent": UserAgent().random(),
        "sec-ch-ua": "^\\^Google",
        "sec-ch-ua-mobile": "?0",
        "sec-ch-ua-platform": "^\\^Windows^^",
    }
    url = "https://www.cnnvd.org.cn/web/cnnvdVul/getCnnnvdDetailOnDatasource"
    data = '{"id":"%s","vulType":"0","cnnvdCode":"%s"}' % (loophole_id,cnnvd_code)
    # response_data = requests.post(url=url, headers=headers, data=data).text
    # print(response_data)
    response_data = requests.post(url=url, headers=headers, data=data).json().get('data').get('cnnvdDetail')
    content += f"【漏洞名称】{response_data.get('vulName')}\n"
    content += f"【漏洞等级】{hazardLevel[response_data.get('hazardLevel')]}\n"
    content += f"【CNNVD编号】{response_data.get('cnnvdCode')}\n"
    content += f"【CVE编号】{response_data.get('cveCode')}\n"
    content += f"【CVE】{response_data.get('cveCode')}\n"
    content += f"【漏洞类型】{response_data.get('vulType')}\n"
    content += f"【漏洞描述】{response_data.get('vulDesc')}\n"
    content += f"【收录时间】{response_data.get('publishTime')}\n"
    content += f"【更新时间】\n{response_data.get('updateTime')}\n"
    content += f"【官方补丁】{response_data.get('patch')}\n"
    content += f"【参考网址】\n{response_data.get('referUrl')}\n"
    content += f"----------\n"
    """
    {'code': 200,
     'data': {'cnnvdDetail': {'affectedProduct': None,
                              'affectedSystem': None,
                              'affectedVendor': 'Subrion',
                              'cnnvdCode': 'CNNVD-202311-299',
                              'cnnvdFiledShow': 'vul_name,cnnvd_code,_code,publish_time,is_official,vendor,hazard_level,vul_type,vul_desc,affected_product,affected_vendor,product_desc,affected_system,refer_url,patch_id,product,update_time,patch',
                              'createTime': None,
                              'createUid': None,
                              'createUname': None,
                              'cveCode': 'CVE-2023-46947',
                              'cveFiledShow': None,
                              'cveVulVO': None,
                              'deleted': None,
                              'hazardLevel': 1,
                              'huaweiFiledShow': None,
                              'huaweiVulVO': None,
                              'ibmFiledShow': None,
                              'ibmVulVO': None,
                              'icsCertFiledShow': None,
                              'icsCertVulVO': None,
                              'id': None,
                              'isOfficial': 0,
                              'microsoftFiledShow': None,
                              'microsoftVulVO': None,
                              'nvdFiledShow': None,
                              'nvdVulVO': None,
                              'patch': None,
                              'patchId': None,
                              'productDesc': None,
                              'publishTime': '2023-11-03 00:00:00',
                              'referUrl': '来源:MISC\r\n'
                                          '链接:https://github.com/intelliants/subrion/issues/909\r\n'
                                          '\r\n'
                                          '来源:cxsecurity.com\r\n'
                                          '链接:https://cxsecurity.com/cveshow/CVE-2023-46947/',
                              'updateTime': '2023-11-13 00:00:00',
                              'updateUid': None,
                              'updateUname': None,
                              'varchar1': '其他',
                              'vendor': '1005853',
                              'version': None,
                              'vulDesc': 'Subrion '
                                         'CMS是Subrion团队的一套基于PHP的内容管理系统（CMS）。该系统可被集成到网站，并支持多种扩展插件等。\r\n'
                                         'Subrion CMS '
                                         '4.2.1版本存在安全漏洞。攻击者利用该漏洞可以远程执行代码。',
                              'vulName': 'Subrion CMS 安全漏洞',
                              'vulType': '其他',
                              'vulTypeName': '其他'},
              'receviceVulDetail': None},
     'message': '操作成功',
     'success': True,
     'time': '2023-11-20 17:05:15'}
    
    """
    pprint.pprint(response_data)
    return content



def _send_info(e_context: EventContext, content: str):
    reply = Reply(ReplyType.TEXT, content)
    channel = e_context["channel"]
    channel.send(reply, e_context["context"])


if __name__ == '__main__':
    # content = ""
    # number = "CNNVD-202311-224"
    # result = search_cnnvd(number=number)
    # if isinstance(result,str):
    #     content += result
    # print(content)
    print(get_loophole(loophole_id="c03637953d3b498f8f0becc3b6989e35"))


